package main

import (
	"crypto/sha256"
	"encoding/json"
	"errors"
	"io"
	"os"
	"path/filepath"
	"reflect"
)

var PUB = [32]byte{
	0x01, 0x06, 0xa5, 0x5a, 0x79, 0xae, 0x1f, 0x7c,
	0x02, 0x06, 0xa5, 0x5a, 0x79, 0xae, 0x1f, 0x7c,
	0x03, 0x06, 0xa5, 0x5a, 0x79, 0xae, 0x1f, 0x7c,
	0x04, 0x06, 0xa5, 0x5a, 0x79, 0xae, 0x1f, 0x7c,
}

type header struct {
	Magic   uint64   `json:"magic"`
	Length  int64    `json:"length"`
	Sha256  [32]byte `json:"sha256"`
	RawName string   `json: "rawname"`
	nb      int
	np      int
	pub     []byte
	reader  io.Reader
}

func (h *header) Read(p []byte) (n int, err error) {
	h.reader.Read(p)
	h.np = h.nb
	for i, _ := range p {
		p[i] ^= h.pub[h.np%32]
		h.np += 1
	}
	return int(h.np - h.nb), nil
}

func en_de_cryption(data []byte) []byte {
	for i, _ := range data {
		data[i] ^= PUB[i%32]
	}

	return data
}

func file_en_cryption(path string, key string) ([]byte, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	header := &header{Magic: 0xbadbeef0, pub: []byte(key)}
	header.Sha256 = sha256.Sum256([]byte(key))
	header.RawName = filepath.Base(path)
	if fileinfo, err := f.Stat(); err == nil {
		header.Length = fileinfo.Size()
	} else {
		return nil, err
	}

	tempFile := os.File{}
	defer func() {
		tempFile.Close()
	}()
	if File, err := os.CreateTemp("", path+".tscz"); err == nil {
		defer func() {
			File.Close()
		}()

		if data, err := json.Marshal(header); err == nil {
			metadataLength := len(data)
			metadataBuf := make([]byte, 4)
			metadataBuf[0] = byte(metadataLength & 0xff)
			metadataBuf[1] = byte(metadataLength >> 8 & 0xff)
			metadataBuf[2] = byte(metadataLength >> 16 & 0xff)
			metadataBuf[3] = byte(metadataLength >> 24 & 0xff)
			File.Write(metadataBuf)

			buffer := make([]byte, 512-4)
			copy(buffer, data)
			File.Write(buffer)
		}

		s := make([]byte, 4096)
		for {
			if nr, err := f.Read(s[:]); err == nil && nr > 0 {

				// 打印原始数据
				//fmt.Println(string(s))
				s = s[:nr]
				header.np = header.nb
				for i, _ := range s {
					s[i] ^= header.pub[header.np%len(header.pub)]
					header.np += 1
				}
				File.Write(s)

			} else {
				break
			}
		}

		os.Rename(File.Name(), path+".tscz")
	}

	return nil, nil
}

func SliceEquals(a, b interface{}) bool {
	_a := reflect.ValueOf(a)
	_b := reflect.ValueOf(b)

	if _a.Len() != _b.Len() {
		return false
	}

	for i := 0; i < _a.Len(); i++ {
		if _a.Index(i).Interface() != _b.Index(i).Interface() {
			return false
		}
	}
	return true
}

func file_de_cryption(path string, key string) ([]byte, error) {
	if filepath.Ext(path) != ".tscz" {
		return nil, errors.New("ext name of de cryption should be .tscz")
	}

	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	header := &header{pub: []byte(key)}

	buf := make([]byte, 512)
	if n, err := f.Read(buf); err == nil && n > 0 {
		metadataLength := buf[0] | buf[1]<<8 | buf[2]<<16 | buf[3]<<24
		if err := json.Unmarshal(buf[4:4+metadataLength], &header); err != nil {
			return nil, err
		}
	} else {
		return nil, err
	}

	if header.Magic != 0xbadbeef0 {
		return nil, errors.New("unidentified by this tool")
	}

	if !SliceEquals(header.Sha256, sha256.Sum256([]byte(key))) {
		return nil, errors.New("invalid password")
	}

	tempFile := os.File{}
	defer func() {
		tempFile.Close()
	}()
	if File, err := os.CreateTemp("", path); err == nil {
		defer func() {
			File.Close()
		}()
		f.Seek(512, os.SEEK_SET)

		s := make([]byte, 4096)
		for {
			if nr, err := f.Read(s[:]); err == nil && nr > 0 {

				s = s[:nr]
				header.np = header.nb
				for i, _ := range s {
					s[i] ^= header.pub[header.np%len(header.pub)]
					header.np += 1
				}
				File.Write(s)
				// 打印解密后的数据
				//fmt.Println(string(s))

			} else {
				break
			}
		}

		os.Rename(File.Name(), header.RawName+".bak")
	}

	return nil, nil
}

func main() {
	// 数据加密、解密
	//data := []byte("hello world")

	//fmt.Printf("%s\n", data)
	//fmt.Printf("%s\n", en_de_cryption(data))
	//fmt.Printf("%s\n", en_de_cryption(data))

	// 文件加密、解密
	file_en_cryption("file.txt", "nihao")

	file_de_cryption("file.txt.tscz", "nihao")
}
